'''
Created on Aug 2, 2010

@author: oabalbin
'''

import numpy as np
from optparse import OptionParser

class mybipartite:
    def __init__(self,adjMat,rownames, colnames):
        """
        B is a matrix of incidence for a bipartite network
        rownames = dict sample name: row location, colnames dict gene,col
        """
        self.B, self.rownames, self.colnames = adjMat, rownames, colnames
    
    def project_samples(self):
        return np.dot(self.B, np.transpose(self.B))
    
    def project_genes(self):
        return np.dot(np.transpose(self.B),self.B)
    

def get_headers(filename):
    inputfile = open(filename)
    first,genedict, sampledict = 0, {}, {}
    genes = inputfile.next().strip("\n").split("\t")
    genes.pop(first)
    for i,g in enumerate(genes):
        genedict[i]=g
    
    for j,line in enumerate(inputfile):
        sample=line.strip("\n").split("\t")[0]
        sampledict[j]=sample
    
    inputfile.close()
    return genedict,sampledict

def read_outlier_matrix(filename):
    
    genedict,sampledict = get_headers(filename)
    nheaders =1
    thiscols = np.array(sorted(genedict.keys())) + 1
    #print thiscols
    # B is a matrix of incidence in a bipartite network
    
    B = np.loadtxt(filename,comments='#', delimiter="\t", skiprows=nheaders, 
                  usecols=thiscols )
    return mybipartite(B,sampledict,genedict)
    
    
def write_projection(P, nodenames, nodes_fname, edges_fname):
    """
    A Projection matrix
    """
    nodesfile, edgesfile=open(nodes_fname,'w'), open(edges_fname,'w')
    Pu = np.triu(P)
    
    nodesfile.write("node"+'\t'+"weight"+'\n')
    eheader = ["source","target","interaction","floating point attribute"]
    edgesfile.write(",".join(eheader).replace(',','\t')+'\n')
    
    for i in range(Pu.shape[0]):
        k=i+1
        node_weight = Pu[i,i]
        edges_weight = list(Pu[i,k:])
        node_a = nodenames[i]
        
        nodeline = [node_a, node_weight]
        nodesfile.write(",".join(map(str,nodeline)).replace(',','\t')+'\n')
        
        # source  target  interaction  boolean attribute  string attribute   
        #     floating point attribute
        for j in range(len(edges_weight)):
            node_b=nodenames[k+j]            
            if edges_weight[j] == 0:
                continue
            
            eline = [node_a, node_b, "pp", edges_weight[j]]
            #if node_a=="SPINK1":
            #    print i, k, j, len(edges_weight), eline
            edgesfile.write(",".join(map(str,eline)).replace(',','\t')+'\n')


def write_bipartite(Pu, nodenames, colnames, nodes_fname, edges_fname):
    """
    A Projection matrix
    """
    nodesfile, edgesfile=open(nodes_fname,'w'), open(edges_fname,'w')
    #Pu = np.triu(P)
    
    nodesfile.write("node"+'\t'+"group"+'\n')
    eheader = ["source","target","interaction","floating point attribute"]
    edgesfile.write(",".join(eheader).replace(',','\t')+'\n')
    
    for i in range(Pu.shape[0]):
        #k=i+1
        #node_weight = Pu[i,i]
        edges_weight = list(Pu[i,:])
        node_a = nodenames[i]
        
        nodeline = [node_a, "sample"]
        nodesfile.write(",".join(map(str,nodeline)).replace(',','\t')+'\n')
        
        # source  target  interaction  boolean attribute  string attribute   
        #     floating point attribute
        for j in range(len(edges_weight)):
            node_b=colnames[j]            
            if edges_weight[j] == 0:
                continue
            
            eline = [node_a, node_b, "pp", edges_weight[j]]
            #if node_a=="SPINK1":
            #    print i, k, j, len(edges_weight), eline
            edgesfile.write(",".join(map(str,eline)).replace(',','\t')+'\n')
    
    
    for j in range(Pu.shape[1]):
        gline = [colnames[j],"gene"]
        nodesfile.write(",".join(map(str,gline)).replace(',','\t')+'\n')
        

    


if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--matrixfile", dest="matrixfile",
                            help="mygenset file to use")
   
    (options, args) = optionparser.parse_args()

    filename = options.matrixfile
        
    myBi = read_outlier_matrix(filename)
    myBi_g = myBi.project_genes()
    myBi_s = myBi.project_samples()
    
    nodes_fname = options.matrixfile+"_projection_samples_nodes.dat"
    edges_fname = options.matrixfile+"_projection_samples_edges.dat"
    
    write_projection(myBi_s, myBi.rownames, nodes_fname, edges_fname)
    
    nodes_fname = options.matrixfile+"_projection_genes_nodes.dat"
    edges_fname = options.matrixfile+"_projection_genes_edges.dat"
    
    write_projection(myBi_g, myBi.colnames, nodes_fname, edges_fname)
    #write_bipartite(myBi.B, myBi.rownames, myBi.colnames, nodes_fname, edges_fname)
    

